import math

import torch
import torch.nn as nn
import torch.optim as optim


class GRU(nn.Module):
    def __init__(self, input_num, hidden_num, output_num):
        super(GRU, self).__init__()

        # 这里设置了 batch_first=True, 所以应该 inputs = inputs.view(inputs.shape[0], -1, inputs.shape[1])
        # 针对时间序列预测问题，相当于将时间步（seq_len）设置为 1。
        self.GRU = nn.GRU(
            input_size=input_num,
            hidden_size=hidden_num,
            batch_first=True
        )

        self.output_linear = nn.Linear(hidden_num, output_num)
        self.hidden = None

    def forward(self, x):
        # h_n of shape (num_layers * num_directions, batch, hidden_size)
        # 这里不用显式地传入隐层状态 self.hidden
        x1 = x.permute(1, 0, 2).contiguous()  # x1: [window, batch, n_val]
        _, h = self.GRU(x1)  # r: [1, batch, hidRNN]
        h = torch.squeeze(h, 0)  # r: [batch, hidRNN]
        res = self.output_linear(h)  # res: [batch, n_val]
        return res


def train(data, x, y, model, criterion, optim, batch_size):
    assert isinstance(model, nn.Module)
    model.train()
    total_loss = n_samples = 0
    for x_data, y_data in data.get_batches(x, y, batch_size, True):
        model.zero_grad()
        output = model(x_data)
        if y_data.size(0) != output.size(0):
            break
        scale = data.scale.expand(output.size(0), data.m)
        loss = criterion(output * scale, y_data * scale)
        loss.backward()
        optim.step()
        total_loss += loss.item()
        n_samples += (output.size(0) * data.m)
    return total_loss / n_samples


def evaluate(data, x, y, model, evaluate_l2, evaluate_l1, batch_size):
    assert isinstance(model, nn.Module)
    model.eval()
    total_loss = total_loss_l1 = n_samples = 0
    test = predict = None
    for x_data, y_data in data.get_batches(x, y, batch_size, False):
        output = model(x_data)
        if y_data.size(0) != output.size(0):
            break
        if predict is None:
            predict = output
            test = y_data
        else:
            predict = torch.cat((predict, output))
            test = torch.cat((test, y_data))
        scale = data.scale.expand(output.size(0), data.m)
        total_loss += evaluate_l2(output * scale, y_data * scale).item()
        total_loss_l1 += evaluate_l1(output * scale, y_data * scale).item()
        n_samples += (output.size(0) * data.m)
    rse = math.sqrt(total_loss_l1 / n_samples) / data.rse
    rae = (total_loss_l1 / n_samples) / data.rae

    predict = predict.data.cpu().numpy()
    y_test = test.data.cpu().numpy()
    sigma_p = predict.std(axis=0)
    sigma_g = y_test.std(axis=0)
    mean_p = predict.mean(axis=0)
    mean_g = y_test.mean(axis=0)
    index = (sigma_g != 0)
    correlation = ((predict - mean_p).T * (y_test - mean_g)).mean(axis=0) / (sigma_p * sigma_g)
    correlation = correlation[index].mean()
    return rse, rae, correlation, predict


def main():
    model = GRU(1, 4, 1)
    print(sum([p.nelement() for p in model.parameters()]))

    criterion = nn.MSELoss(reduction='sum').cuda()
    evaluate_l1 = nn.L1Loss(reduction='sum').cuda()
    evaluate_l2 = nn.MSELoss(reduction='sum').cuda()

    batch_size = 168
    epochs = 100
    best_value = 10000000
    save = 'GRU.pt'
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

    import dataset
    for epoch in range(1, epochs):
        train_loss = train(
            dataset.data, dataset.data.train[0], dataset.data.train[1],
            model, criterion, optimizer, batch_size)

        rse, rae, corr, predict = evaluate(
            dataset.data, dataset.data.train[0], dataset.data.train[1],
            model, evaluate_l2, evaluate_l1, batch_size)

        print(train_loss, rse, rae, corr, predict)

        if rse < best_value:
            with open(save, 'wb') as f:
                torch.save(model, f)
            best_value = rse

    # truth = dataset.data.test[1].numpy()
    # for i in range(truth.shape[1]):
    #     plt.figure(figsize=(10, 5))
    #     plt.plot(truth[:24 * 7 * 4, i], label='ground truth')
    #     plt.plot(pred[:24 * 7 * 4, i], label='prediction')
    #     plt.legend()
    #     plt.show()


if __name__ == "__main__":
    main()
